import json
import json_repair
import pandas as pd
from tqdm import tqdm, trange
from kdbUtil import * 
import uuid
import os

from fastapi.responses import StreamingResponse
from duckduckgo_search import DDGS
import embeddingDB
from embeddingDB import embed_with_list_of_str,pgDB
from fastapi import APIRouter, Query, UploadFile, File,Form,Body
router = APIRouter(
    prefix= "/aiSearch",
    tags= ["aiSearchRouter"],
    responses= {404: {"description": "Not found path knowledge"}}
)

class search(BaseModel):
    keywords: Optional[str] = ""
    kdbListChecked: Optional[list] = []
    
# vdb = VectorDB()
# docs = splitDocument("./d1.docx")
# docs = {"documents":docs,"metadatas":[{"id":uuid.uuid4().hex} for i in docs],"ids":[uuid.uuid4().hex for i in docs]}
# vdb.persistentEmbedding(docs)
# userQuestion = "企业大脑"
# localKnowledgeRS = vdb.query(userQuestion,10)
# internetRS = internetSearch(userQuestion,10)
# rs = {"localKnowledgeRS":localKnowledgeRS,"internetRS":internetRS}
# print(rs)
# print(active_apis_data[0])

# active_apis_data = pd.DataFrame(active_apis_data)
import time
@router.post("/searchDoc")
def searchDoc(params:search):
    rs = {"code":200,"msg":"创建成功","data":{}}
    embeddings = embed_with_list_of_str([params.keywords])[0]
    print(params)
    searchRS = {"kdbRSList":[],"internetRS":[]}
    # searchRS["internetRS"] = internetSearch(params.keywords)
    for kdbTitle in params.kdbListChecked:
        sql = '''
        SELECT content,create_time,file_name,id,kdb_id,download_url FROM knowledge
        ORDER BY vector <-> '{}' LIMIT 10;
        '''.format(embeddings)
        db = pgDB()
        kdbRS = db.select(sql)
        db.close()
        # kdb = embeddingDB.VectorDB()
        # [k.update({"download_url":"http://tech.zscampus.com/file/uploadFiles/{}/{}".format(k["kdb_id"],k["file_name"]) }) for k in kdbRS]
        # kdb.getCollection(kdbTitle)
        # kdbRS = kdb.query(params.keywords,10)
        for k in kdbRS:
            k["download_url"] = "zsai-algorithm/{}".format(k["download_url"]) 
        searchRS["kdbRSList"].append({"kdbTitle":kdbTitle,"kdbRS":kdbRS})
    # kdbList = [embeddingDB.VectorDB() for kdb in params.kdbListChecked]
    # [kdb.getCollection() for kdb in kdbList]
    time.sleep(1)
    rs["data"]=searchRS
    return rs
    
from fastapi.concurrency import run_in_threadpool
@router.post("/answer")
async def answer(params: search):
    rs = {"code": 200, "msg": "创建成功", "data": {}}
    print(params)
    
    searchRS = {"kdbRSList": [], "internetRS": []}
    rsList = []
    embeddings = embed_with_list_of_str([params.keywords])[0]
    
    for kdbTitle in params.kdbListChecked:
        sql = '''
        SELECT * FROM knowledge ORDER BY vector <-> '{}' LIMIT 10;
        '''.format(embeddings)
        db = pgDB()
        kdbRS = await run_in_threadpool(db.select, sql)  # 确保这里使用了 await
        db.close()
        [rsList.append(k["content"]) for k in kdbRS]
        searchRS["kdbRSList"].append({"kdbTitle": kdbTitle, "kdbRS": kdbRS})
    
    rs["data"] = rsList
    
    prompt = '''
    <你的角色>
    你是智能搜索解题专家，能够根据但不完全依靠我提供的背景知识，解答我的问题。
    
    <背景知识>
    格式为 序号.知识，序号用于你做引用。
    {}
    
    <你的任务>
    1. 请一步步思考并深入回答。用MarkDown格式条理清晰的回答。
    2. 若你的输出参考了我提供的某些背景知识，则请在你的输出中，增加引用，引用来源于你参考的背景知识的序号。
    示例：如某个回答内容引用了背景知识中序号为x的知识，则再内容后追加[x]。
    
    
    <我的问题>
    {}
    
    '''.format("\n".join(["> {}.{}".format(i + 1, v) for i, v in enumerate(rsList)]), params.keywords)
    
    ak = 'fastgpt-xw2742EViMJdNNBNgWgYJwMFKxZWtREdiyvY45v2dA7MBXsHGlzoAOe1'
    
    return StreamingResponse(hanguang_agent(prompt, ak), media_type="text/event-stream")
    

def internetSearch(question,nums=10):
    r=[]
    try:
        with DDGS(proxy="http://localhost:26004", timeout=20) as ddgs:
            r= ddgs.text(question, max_results=nums)
            # print(r)
    except Exception as e:
        print(e)
        pass
    return r

import requests
import random
from json_repair import repair_json 
import httpx
async def hanguang_agent(content, apikey):
    url = 'https://model.zscampus.com/api/v1/chat/completions'
    headers = {
        'Authorization': f'Bearer {apikey}',
        'Content-Type': 'application/json'
    }
    data = {
        "chatId": uuid.uuid4().hex,
        "stream": True,
        "detail": False,
        "messages": [
            {
                "content": content,
                "role": "user"
            }
        ]
    }

    async with httpx.AsyncClient() as client:
        async with client.stream("POST", url, json=data, headers=headers) as response:
            async for chunk in response.aiter_text():
                if "[DONE]" not in chunk:
                    try:
                        chunk_data = json.loads(chunk[5:].replace(' ', ''))
                        yield chunk_data["choices"][0]["delta"]["content"]
                    except json.JSONDecodeError:
                        continue
                else:
                    yield ""